#!/usr/bin/python3
"""
Robot controller

By zen shawn Aug 16th, 2020
"""

import logging
import time
from threading import Thread
from typing import Dict, List

import numpy as np
from vispy.scene import transforms
from vispy.scene.canvas import SceneCanvas
from vispy.scene.visuals import Mesh

from .MeshLoader import LoadMesh, LoadUrdf
from .modern_robotics import *
from .Trajectory import ComputeTrajectory
from .Utils import *
from .XmlParser import RobotTree

logger = logging.getLogger(__file__)


class Robot:
    def __init__(self, interp_period=1/24):
        """Construct a robot

        Keyword Arguments:
            name {str,None} -- robot name (default: {None})
            interp_period {float} -- robot control interpolation period (default: {0.01})
        """
        self.interp_period = interp_period
        self.joint_positions = np.zeros(0)
        self.base_transformation = np.eye(4)
        self.meshes = {}  # Dict(link_name, link_mesh)

    def Show(self, urdf_path: str, canvas: SceneCanvas):
        self.urdf_path = urdf_path
        name, tree, link_meshes = LoadUrdf(
            canvas, self.urdf_path, show_edge=False, show_face=True)
        self.tree = tree
        self.name = name
        self.link_meshes = link_meshes
        self.joint_positions = np.zeros(self.tree.GetDoF())
        self.canvas = canvas
        self.canvas.events.draw.connect(self._Sync)
        joints = self.tree.GetJoints()
        # upper / lower position limits
        jls_lu = np.ones((self.tree.GetDoF(), 2))
        jls = np.ones((3, self.tree.GetDoF()))  # vel, acc, jerk
        for i, joint_name in enumerate(self.tree.active_joints):
            jls[0, i] = joints[joint_name]['limit'].get('velocity', 1.0)
            jls[1, i] = joints[joint_name]['limit'].get('acceleration', 1.0)
            jls[2, i] = joints[joint_name]['limit'].get('jerk', 1.0)
            jls_lu[i, 0] = joints[joint_name]['limit'].get('lower', -1.57)
            jls_lu[i, 1] = joints[joint_name]['limit'].get('upper', 1.57)
        self.joint_limits = jls
        self.vlim = np.c_[-self.joint_limits[0], self.joint_limits[0]].T
        self.alim = np.c_[-self.joint_limits[1], self.joint_limits[1]].T
        self.joint_position_limits = jls_lu

    def Remove(self):
        self.canvas.events.draw.disconnect(self._Sync)
        for link in self.link_meshes.values():
            link.parent = None

    def RandomState(self, wait=False):
        joints = np.random.uniform(
            self.joint_position_limits[:, 0], self.joint_position_limits[:, 1], self.tree.GetDoF())
        return self.MoveJ(joints, wait)

    def FK(self, joints):
        pose = FKinSpace(self.tree.m_home_pose, self.tree.m_s_list, joints)
        return pose

    def IK(self, pose):
        joints = IKinSpace(self.tree.m_s_list, self.tree.m_home_pose,
                           pose, self.GetRobotJointPosition(), 0.01, 0.001)
        return joints

    def GetName(self):
        return self.name

    def SetBaseTransformation(self, pose):
        """Set robot base transformation

        Arguments:
            pose {pose} -- pose data, 4x4 General Transform Matrix

        Returns:
            bool -- result
        """
        self.base_transformation = pose
        self.link_meshes[self.tree.GetRootLink(
        )].transform = transforms.MatrixTransform(NormToVisForm(pose))
        return True

    def GetRobotJointPosition(self):
        return self.joint_positions

    def GetRobotPose(self):
        return self.FK(self.GetRobotJointPosition())

    def MoveJ(self, joints, wait=False):
        """Move robot joints

        Arguments:
            joints {list, np.array} -- target joints, dimension: dof

        Returns:
            bool -- result
        """
        traj = ComputeTrajectory(
            np.array([self.GetRobotJointPosition(), joints]), self.vlim, self.alim)
        return self.FollowTraj(traj, wait)

    def MoveL(self, pose, wait=False):
        joints, ret = self.IK(pose)
        if ret:
            return self.MoveJ(joints)
        else:
            logger.error("IK error at pose:\n{}".format(pose))
            return False

    def FollowTraj(self, traj, wait=False):
        """Follow a toppra trajecotry

        Arguments:
            traj {class} -- joint trajectory, should have property "duration" and "eval" method
            wait (bool, optional): blocking or new thread. Defaults to False.

        Returns:
            bool -- result
        """
        if not wait:
            thread = Thread(target=self.FollowTraj, args=(
                traj,), kwargs={"wait": True}, daemon=False)
            thread.start()
            return True
        t0 = time.time()
        t = 0
        while t < traj.duration:
            self._ServoJ(traj.eval(t))
            t = time.time() - t0
            time.sleep(self.interp_period)
        return True

    def _ServoJ(self, joints):
        if len(joints) != self.tree.GetDoF():
            logger.error(f"input {joints} has different dof!")
            return False
        self.joint_positions = np.asarray(joints, dtype=float)

    def _Sync(self, event):
        """Synchronize jonit position to vispy

        Args:
            event ([type]): [description]

        Returns:
            [type]: [description]
        """
        for i, rad in enumerate(self.joint_positions):
            link = self.link_meshes[self.tree.active_links[i]]
            R = Rotation.from_rotvec(rad * self.tree.active_joints_axis[i])
            T = link.transform.matrix
            # TODO: check is vispy Transform just transpose of general Transform
            T[:3, :3] = R.as_matrix().T
            link.transform = transforms.MatrixTransform(T)
        return True

    def __repr__(self):
        return "<Robot at {:#x}, name: {}>".format(id(self), self.name)
